import torch

device = torch.device("cpu")


def eucli(a, b):
    # 计算欧氏距离
    dist = torch.sqrt(torch.sum((a - b) ** 2))
    return dist


"""
state: 3 * 21 
node : 21
vec: 3
actions: 
reward: 
done: 

"""


class Env():
    # 静态不变信息
    n_samples = 1
    n_customer = 20
    n_vehicle = 3
    CAPACITIES = {10: 20., 20: 30., 50: 40., 100: 50.}

    def __init__(self):
        # 静态不变信息
        self.speed = 1
        self.depot_xy = torch.rand((self.n_samples, 2), device=device)  # 车场坐标
        self.customer_xy = torch.rand((self.n_samples, self.n_customer, 2), device=device)  # 节点坐标
        self.xy = torch.cat((self.depot_xy.unsqueeze(1), self.customer_xy), 1).to(device)
        self.demand = torch.randint(size=(self.n_samples, self.n_customer), low=1, high=10, device=device) / \
                      self.CAPACITIES[self.n_customer]  # 客户节点的需求
        self.capacity = torch.empty(size=(self.n_samples, self.n_vehicle), device=device).uniform_(5.0,
                                                                                                   10.0)  # 3辆车的车辆容量
        # 生成数组Depot，维度为（5, 1, 2）
        array_D = torch.zeros(5, 1, 2)
        array_D[:, :, 0] = 0  # 第一列赋值为0
        array_D[:, :, 1] = 10  # 第二列赋值为5
        # 生成数组Customer，维度为（5, 20, 2）
        array_C = torch.zeros(5, 20, 2)
        array_C[:, :, 0].uniform_(0, 1.0)  # 在[0, 0.1]范围内随机生成第一列的值
        array_C[:, :, 1].uniform_(1.0, 2.0)  # 在[0.1, 0.2]范围内随机生成第二列的值
        # 拼接数组D和数组C在第二维上
        self.time_window = torch.cat((array_D, array_C), dim=1)  # 客户的时间窗口

        # 动态可变信息
        # 每辆车当前剩余容量
        self.cur_cap = self.capacity.clone()
        # 每辆车的当前时间
        self.cur_time = torch.zeros(self.n_samples, self.n_vehicle)
        # 掩码,初始depot不可访问
        self.mask = torch.zeros((self.n_samples, self.n_vehicle, self.n_customer + 1))
        self.mask[:, :, 0] = 1
        # 用车辆的索引来标记每个节点是哪个车辆来送货，初始化都为100
        # self.flag = [[] for _ in range(self.n_samples)]#100 * torch.ones((self.n_samples, self.n_customer))
        # 已访问节点
        self.all_venodes = [[[] for _ in range(self.n_vehicle)] for _ in range(self.n_samples)]
        self.tours = torch.zeros([self.n_samples, self.n_vehicle, (self.n_customer + 1) * 2], dtype=torch.int32).to(
            device)  # [5,3,42]
        # 目标函数值
        self.all_cost = torch.zeros(self.n_samples, 1)
        # 返回节点次数
        # self.re_depot=0
        # 车辆是否返回网点，未返回为0，返回为1
        # self.ve_redepot = torch.zeros((self.n_samples, self.n_vehicle))

    """
        环境重置
    """

    def reset(self):
        # 每辆车当前剩余容量
        self.cur_cap = self.capacity.clone()
        # 每辆车的当前时间
        self.cur_time = torch.zeros(self.n_samples, self.n_vehicle)
        # 掩码,初始depot不可访问
        self.mask = torch.zeros((self.n_samples, self.n_vehicle, self.n_customer + 1))
        self.mask[:, :, 0] = 1
        # 用车辆的索引来标记每个节点是哪个车辆来送货，初始化都为100
        # self.flag = [[] for _ in range(self.n_samples)] #100 * torch.ones((self.n_samples, self.n_customer))
        # 已访问节点
        self.all_venodes = [[[] for _ in range(self.n_vehicle)] for _ in range(self.n_samples)]
        self.tours = torch.zeros([self.n_samples, self.n_vehicle, (self.n_customer + 1) * 2], dtype=torch.int32).to(
            device)  # [5,3,42]
        # 目标函数值
        self.all_cost = torch.zeros(self.n_samples, 1).to(device)
        # 返回节点次数
        # self.re_depot = torch.zeros(self.n_samples, 1)
        # 车辆是否返回网点，未返回为0，返回为1
        # self.ve_redepot = torch.zeros((self.n_samples, self.n_vehicle))
        # cur_cap:容量， mask: node使用情况
        return self.mask, self.cur_cap

    def so_mask(self, all_venodes):  # 回归初始状态,坐标轴原点在左上角（5，n）
        # 拼接得到的二维列表
        done = False
        all_nodes = [sum(sublist, []) for sublist in all_venodes]
        self.mask = torch.zeros((self.n_samples, self.n_vehicle, self.n_customer + 1))
        self.mask[:, :, 0] = 1

        # 所有节点都访问完了，并且所有车辆都返回车厂
        if len(all_nodes[0]) >= self.n_customer:  # and (self.ve_redepot == 1).all(dim=1).all().item():
            self.mask = torch.ones((self.n_samples, self.n_vehicle, self.n_customer + 1))
            done = True
        # 判断是否所有节点都访问完了,但是所有节点还未返回车场
        # elif len(all_nodes[1])>=self.n_customer:
        #     self.mask = torch.ones((self.n_samples, self.n_vehicle, self.n_customer + 1))
        #     # self.mask[:, :, 0] = 0
        #     #给没回车场的车辆掩码为0
        #     for i in range(self.n_samples):
        #         for j in range(self.n_vehicle):
        #             if self.ve_redepot[i][j]==0:
        #                 self.mask[i][j][0] = 0
        # 还未访问完
        else:
            for row in range(self.n_samples):
                for sel_nodes in all_nodes[row]:
                    self.mask[row, :, sel_nodes] = 1
                # 每辆车掩码
                for k in range(self.n_vehicle):
                    # 判断客户从下标1开始
                    for j in range(1, self.n_customer + 1):
                        if self.mask[row][k][j] == 0 and self.cur_cap[row][k] < self.demand[row][j - 1]:
                            self.mask[row][k][j] = 1
        return self.mask, done

    def step(self, selected_vehicle, current_node):  # 外部调用这个函数来改变当前位置
        # selected_vehicle(5,1)当前选择车辆，current_node(5,1)当前选择节点
        # print(selected_vehicle)
        for ep in range(self.n_samples):
            cur_ve = selected_vehicle[ep][0]
            cur_node = current_node[ep][0]
            # 如果返回车场
            # if cur_node == 0:
            #     # self.re_depot[ep] += 1
            #     self.ve_redepot[ep][cur_ve]=1
            # 上一个节点
            # aaa=len(self.all_venodes[ep][cur_ve])
            if len(self.all_venodes[ep][cur_ve]) == 0:
                pre_node = 0
            else:
                pre_node = self.all_venodes[ep][cur_ve][-1]
            # 距离
            distance = eucli(self.xy[ep][pre_node], self.xy[ep][cur_node])
            self.tours[ep][cur_ve][len(self.all_venodes[ep][cur_ve])] = cur_node
            self.all_venodes[ep][cur_ve].append(cur_node)
            # self.flag[ep].append(cur_ve)
            # 如果不是返回车场，车辆的剩余容量
            if cur_node != 0:
                self.cur_cap[ep][cur_ve] = self.cur_cap[ep][cur_ve] - self.demand[ep][cur_node - 1]
            # 当前时间
            self.cur_time[ep][cur_ve] = self.cur_time[ep][cur_ve] + distance / self.speed
            reward = distance / self.speed
            if self.cur_time[ep][cur_ve] < self.time_window[ep][cur_node][0]:
                reward += (self.time_window[ep][cur_node][0] - self.cur_time[ep][cur_ve]) * 0.5
            elif self.cur_time[ep][cur_ve] > self.time_window[ep][cur_node][1]:
                reward += (self.cur_time[ep][cur_ve] - self.time_window[ep][cur_node][1]) * 2.0
            # 目标函数值
            self.all_cost[ep] += reward
        mask, done = self.so_mask(self.all_venodes)

        return mask, reward, done, self.cur_cap


if __name__ == "__main__":
    env = Env()
    noooood = [[[19, 4], [2, 16], [9, 14]], [[10, 5], [6, 19], [4, 11]], [[1, 9], [15, 6], [7, 4]],
               [[12, 8], [7, 1], [17, 4]], [[7, 15], [16, 9], [14, 10]]]
    # mask=env.so_mask(noooood)
    selected_vehicle = torch.tensor([[2], [0], [2], [0], [1]])
    current_node = torch.tensor([[15], [8], [16], [2], [13]])
    mask, _, _, cur_cap = env.step(selected_vehicle, current_node)
    aaa = 1
